package Others;

import java.util.ArrayList;

public class UnionFindSet {
    ArrayList<Integer> array;

    public UnionFindSet(int nInitialElements) {
        int capacity = 1;
        while (capacity < nInitialElements) capacity *= 2;
        array = new ArrayList<>(capacity);
        for (int i = 0; i < nInitialElements; i++) array.add(-1);
    }

    public int find(int elem) {
        if (elem < 0 || elem >= array.size())
            throw new IndexOutOfBoundsException("Query element " + elem + " is out of bounds.");

        if (array.get(elem) < 0) return elem;
        else {
            int parent = array.get(elem);
            array.set(elem, find(parent)); // 路径压缩
            return array.get(elem);
        }
    }

    public int union(int elem1, int elem2) {
        int root1 = find(elem1), root2 = find(elem2);
        if (root1 == root2) return root1;

        if (array.get(root1) < array.get(root2)) {
            int array2Size = -array.get(root2);
            array.set(root2, root1);
            array.set(root1, array.get(root1) - array2Size);
            return root1;
        }
        else {
            int array1Size = -array.get(root1);
            array.set(root1, root2);
            array.set(root2, array.get(root2) - array1Size);
            return root2;
        }
    }
}
